package com.self.av;

import static android.os.Environment.DIRECTORY_PICTURES;
import static android.os.Environment.getExternalStoragePublicDirectory;

import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.ImageFormat;
import android.graphics.PointF;
import android.graphics.Rect;
import android.graphics.YuvImage;
import android.icu.text.SimpleDateFormat;
import android.os.Bundle;
import android.os.Handler;
import android.os.Looper;
import android.view.View;
import android.widget.Button;
import android.widget.ImageView;
import android.widget.Toast;

import androidx.annotation.NonNull;
import androidx.appcompat.app.AppCompatActivity;
import androidx.camera.core.Camera;
import androidx.camera.core.CameraSelector;
import androidx.camera.core.ImageAnalysis;
import androidx.camera.core.ImageProxy;
import androidx.camera.core.Preview;
import androidx.camera.lifecycle.ProcessCameraProvider;
import androidx.camera.view.PreviewView;
import androidx.core.content.ContextCompat;
import androidx.lifecycle.LifecycleOwner;

import com.google.common.util.concurrent.ListenableFuture;
import com.google.mediapipe.framework.image.BitmapImageBuilder;
import com.google.mediapipe.framework.image.MPImage;
import com.google.mediapipe.tasks.components.containers.NormalizedLandmark;
import com.google.mediapipe.tasks.core.Delegate;
import com.google.mediapipe.tasks.core.OutputHandler;
import com.google.mediapipe.tasks.vision.core.RunningMode;
import com.google.mediapipe.tasks.vision.poselandmarker.PoseLandmarker;
import com.google.mediapipe.tasks.vision.poselandmarker.PoseLandmarkerResult;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class CameraAnalysisActivity extends AppCompatActivity {
    private PreviewView previewView;
    private Button lockPositionButton;
    private Button hideLandmarkerButton;
    private Button hideGridButton;
    private Button screenShotButton;
    private final Handler handler = new Handler(Looper.getMainLooper());
    private com.self.av.GridOverlayView gridOverlayView;
    private ActionLineView actionLineView;
    private ImageView freezeImageView;
    private ImageView screenShotImageView;
    private PoseLandmarker poseLandmarker;
    private ExecutorService cameraExecutor;
    private Bitmap frameBitmap;

    private static final int[] PAIR_IDS = {
            0, 2, 5, 7, 8,
            11, 12, 13, 14,
            15, 16, 23, 24,
            25, 26, 27, 28,
    };

    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_camera_analysis);

        loadMPModel();
        bindCameraUseCases();

        previewView = findViewById(R.id.previewView);
        actionLineView = findViewById(R.id.actionLineView);
        gridOverlayView = findViewById(R.id.gridOverlayView);
        freezeImageView = findViewById(R.id.freezeImageView);
        screenShotImageView = findViewById(R.id.screenShotImageView);
        lockPositionButton = findViewById(R.id.lockPositionButton);
        hideLandmarkerButton = findViewById(R.id.hideLandmarkerButton);
        hideGridButton = findViewById(R.id.hideGridButton);
        screenShotButton = findViewById(R.id.screenShotButton);

        hideLandmarkerButton.setOnClickListener(v -> {
            if (actionLineView.getVisibility() == View.VISIBLE) {
                actionLineView.setVisibility(View.INVISIBLE);
                hideLandmarkerButton.setText("显示标记");
            } else {
                actionLineView.setVisibility(View.VISIBLE);
                hideLandmarkerButton.setText("隐藏标记");
            }
        });

        hideGridButton.setOnClickListener(v -> {
            if (gridOverlayView.getVisibility() == View.VISIBLE) {
                gridOverlayView.setVisibility(View.INVISIBLE);
                hideGridButton.setText("显示网格");
            } else {
                gridOverlayView.setVisibility(View.VISIBLE);
                hideGridButton.setText("隐藏网格");
            }
        });

        lockPositionButton.setOnClickListener(new View.OnClickListener() {
            @Override
            public void onClick(View v) {
                if (lockPositionButton.getText().toString().equals("锁定")) {
//                    Bitmap frameBitmap = getFrameBitmap(player.getCurrentPosition() * 1000);
                    freezeImageView.setImageBitmap(scaledBitmap(frameBitmap));
                    freezeImageView.setVisibility(View.VISIBLE);
                    lockPositionButton.setText("解锁");
//                    Log.e("@@@@@lockPositionButton", "textureView: " + textureView.getWidth()+" "+textureView.getHeight());
                } else {
                    freezeImageView.setImageBitmap(null);
                    freezeImageView.setVisibility(View.INVISIBLE);
                    lockPositionButton.setText("锁定");
                }
            }
        });

        screenShotButton.setOnClickListener(new View.OnClickListener() {
            @Override
            public void onClick(View v) {
                screenShotImageView.setImageBitmap(scaledBitmap(frameBitmap));
                screenShotImageView.setVisibility(View.VISIBLE);

                // 获取当前 Activity 的根视图
                View rootView = getWindow().getDecorView().getRootView();
                rootView.setDrawingCacheEnabled(true);
                Bitmap bitmap = Bitmap.createBitmap(rootView.getDrawingCache());
                rootView.setDrawingCacheEnabled(false);
                String timeStamp = new SimpleDateFormat("yyyyMMdd_HHmmss", Locale.getDefault()).format(new Date());
                String fileName = "AV_" + timeStamp + ".jpg";
                File file = new File(getExternalStoragePublicDirectory(DIRECTORY_PICTURES), fileName);
                try (FileOutputStream fos = new FileOutputStream(file)) {
                    bitmap.compress(Bitmap.CompressFormat.JPEG, 100, fos);
                    fos.flush();
                } catch (IOException e) {
                    e.printStackTrace();
                }
                screenShotImageView.setImageBitmap(null);
                screenShotImageView.setVisibility(View.INVISIBLE);
                Toast.makeText(getApplicationContext(), file.getAbsolutePath(), Toast.LENGTH_SHORT).show();
            }
        });

        cameraExecutor = Executors.newSingleThreadExecutor();
    }

    private void loadMPModel() {
        // 模型文件在assets目录下的文件名
        String modelName = "pose_landmarker_heavy.task";
        com.google.mediapipe.tasks.core.BaseOptions baseOptions = com.google.mediapipe.tasks.core.BaseOptions.builder()
                .setDelegate(Delegate.GPU)
                .setModelAssetPath(modelName)
                .build();

        OutputHandler.ResultListener<PoseLandmarkerResult, MPImage> returnLivestreamResult = (result, input) -> {
            List<List<NormalizedLandmark>> landmarks_lst = result.landmarks();
            for (List<NormalizedLandmark> landmarks : landmarks_lst) {
                List<PointF> jointPoints = convertLandmarkToJointPoints(landmarks);
                List<PointF[]> lines = convertLandmarkToLines(landmarks);
                actionLineView.setCircles(jointPoints);
                actionLineView.setLines(lines);
            }
            actionLineView.invalidate();
        };

        PoseLandmarker.PoseLandmarkerOptions poseLandmarkerOptions = PoseLandmarker.PoseLandmarkerOptions.builder()
                .setBaseOptions(baseOptions)
                .setMinPoseDetectionConfidence(0.5f)
                .setMinTrackingConfidence(0.5f)
                .setMinPosePresenceConfidence(0.5f)
                .setNumPoses(1)
                .setOutputSegmentationMasks(false)
                .setResultListener(returnLivestreamResult)
                .setRunningMode(RunningMode.LIVE_STREAM)
                .build();

        poseLandmarker = PoseLandmarker.createFromOptions(getApplicationContext(), poseLandmarkerOptions);
    }

    private void bindCameraUseCases() {
        ListenableFuture<ProcessCameraProvider> cameraProviderFuture = ProcessCameraProvider.getInstance(this);

        cameraProviderFuture.addListener(() -> {
            try {
                ProcessCameraProvider cameraProvider = cameraProviderFuture.get();
                bindPreview(cameraProvider);
            } catch (ExecutionException | InterruptedException e) {
                // 处理异常
            }
        }, ContextCompat.getMainExecutor(this));
    }

    private void bindPreview(@NonNull ProcessCameraProvider cameraProvider) {
        Preview preview = new Preview.Builder().build();

        CameraSelector cameraSelector = new CameraSelector.Builder()
                .requireLensFacing(CameraSelector.LENS_FACING_BACK)
                .build();

        preview.setSurfaceProvider(previewView.getSurfaceProvider());

        ImageAnalysis imageAnalysis = new ImageAnalysis.Builder()
                .setBackpressureStrategy(ImageAnalysis.STRATEGY_KEEP_ONLY_LATEST)
                .build();

        imageAnalysis.setAnalyzer(cameraExecutor, image -> {
            processImageProxy(image);
//            image.close();
        });

        Camera camera = cameraProvider.bindToLifecycle((LifecycleOwner) this, cameraSelector, preview, imageAnalysis);
    }

    private void processImageProxy(@NonNull ImageProxy imageProxy) {
        frameBitmap = imageProxyToBitmap(imageProxy);
        Bitmap bitmap = Bitmap.createScaledBitmap(frameBitmap, 224, 224, true);
        MPImage mpImage = new BitmapImageBuilder(bitmap).build();
        poseLandmarker.detectAsync(mpImage, System.currentTimeMillis());
    }
    private Bitmap imageProxyToBitmap(ImageProxy image) {
        android.media.Image.Plane[] planes = image.getImage().getPlanes();
        ByteBuffer yBuffer = planes[0].getBuffer();
        ByteBuffer uBuffer = planes[1].getBuffer();
        ByteBuffer vBuffer = planes[2].getBuffer();

        int ySize = yBuffer.remaining();
        int uSize = uBuffer.remaining();
        int vSize = vBuffer.remaining();

        byte[] nv21 = new byte[ySize + uSize + vSize];
        // 将 Y 数据复制到 nv21 数组
        yBuffer.get(nv21, 0, ySize);
        vBuffer.get(nv21, ySize, vSize);
        uBuffer.get(nv21, ySize + vSize, uSize);

        YuvImage yuvImage = new YuvImage(nv21, ImageFormat.NV21, image.getWidth(), image.getHeight(), null);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        yuvImage.compressToJpeg(new Rect(0, 0, yuvImage.getWidth(), yuvImage.getHeight()), 75, out);
        image.close();
        byte[] imageBytes = out.toByteArray();
        Bitmap bitmap = BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.length);
        return bitmap;
    }

    private Bitmap scaledBitmap(Bitmap oriBitmap) {
        int bitmapWidth = oriBitmap.getWidth();
        int bitmapHeight = oriBitmap.getHeight();
        int previewWidth = previewView.getWidth();
        int previewHeight = previewView.getHeight();
        float rateW = (float) previewWidth / bitmapWidth;
        float rateH = (float) previewHeight / bitmapHeight;
        float rate = Math.max(rateW, rateH);
        int newWidth = (int) (bitmapWidth * rate);
        int newHeight = (int) (bitmapHeight * rate);
        Bitmap bitmap = Bitmap.createScaledBitmap(oriBitmap, newWidth, newHeight, true);
        if (rateW > rateH) {
            bitmap = Bitmap.createBitmap(bitmap, 0, (newHeight - previewHeight) / 2, previewWidth, previewHeight);
        } else {
            bitmap = Bitmap.createBitmap(bitmap, (newWidth - previewWidth) / 2, 0, previewWidth, previewHeight);
        }
        return bitmap;
    }

    private List<PointF> convertLandmarkToJointPoints(List<NormalizedLandmark> landmarks) {
        List<PointF> points = new ArrayList<>();
        for (int i : PAIR_IDS) {
            NormalizedLandmark landmark = landmarks.get(i);
            points.add(getJointPosition(landmark));
        }
        return points;
    }

    private List<PointF[]> convertLandmarkToLines(List<NormalizedLandmark> landmarks) {
        List<PointF[]> lines = new ArrayList<>();
        int[][] connections = {
                {0, 1}, {1, 2}, {2, 0},
                {1, 3}, {2, 4},
                {5, 7}, {7, 9}, {6, 8}, {8, 10},
                {5, 6}, {6, 12}, {12, 11}, {5, 11},
                {11, 13}, {13, 15}, {12, 14}, {14, 16}
        };

        PointF[] points = new PointF[17];

        for (int i = 0; i < 17; i++) {
            NormalizedLandmark landmark = landmarks.get(PAIR_IDS[i]);
            points[i] = getJointPosition(landmark);
        }

        for (int[] connection : connections) {
            int startIndex = connection[0];
            int endIndex = connection[1];
            if (points[startIndex] != null && points[endIndex] != null) {
                PointF[] line = new PointF[2];
                line[0] = points[startIndex];
                line[1] = points[endIndex];
                lines.add(line);
            }
        }
        return lines;
    }

    private PointF getJointPosition(NormalizedLandmark landmark) {
        int width = previewView.getWidth();
        int height = previewView.getHeight();
        float rateW = (float) width / frameBitmap.getWidth();
        float rateH = (float) height / frameBitmap.getHeight();
        if (rateW > rateH) {
            float x = landmark.x() * frameBitmap.getWidth() * rateW;
            float y = landmark.y() * frameBitmap.getHeight() * rateW;
            y = y - ((frameBitmap.getHeight() * rateW - height) / 2);
            return new PointF(x, y);
        } else {
            float x = landmark.x() * frameBitmap.getWidth() * rateH;
            float y = landmark.y() * frameBitmap.getHeight() * rateH;
            x = x - ((frameBitmap.getWidth() * rateH - width) / 2);
            return new PointF(x, y);
        }
//        float x = landmark.x() * frameBitmap.getWidth();
//        float y = landmark.y() * frameBitmap.getHeight();
//        return new PointF(x, y);
    }

    @Override
    protected void onDestroy() {
        super.onDestroy();
        cameraExecutor.shutdown();
        if (poseLandmarker != null) {
            poseLandmarker.close();
        }
    }
}